{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Title: #万灵之树"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Difficulty: #Hard"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Category Title: #Algorithms"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Tag Slug: #"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Name Translated: #"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solution Name: treeOfInfiniteSouls"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Title: #万灵之树"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Translated Content:\n",
    "探险家小扣终于来到了万灵之树前，挑战最后的谜题。\n",
    "已知小扣拥有足够数量的链接节点和 `n` 颗幻境宝石，`gem[i]` 表示第 `i` 颗宝石的数值。现在小扣需要使用这些链接节点和宝石组合成一颗二叉树，其组装规则为：\n",
    "- 链接节点将作为二叉树中的非叶子节点，且每个链接节点必须拥有 `2` 个子节点；\n",
    "- 幻境宝石将作为二叉树中的叶子节点，所有的幻境宝石都必须被使用。\n",
    "\n",
    "能量首先进入根节点，而后将按如下规则进行移动和记录：\n",
    "- 若能量首次到达该节点时：\n",
    "    - 记录数字 `1`；\n",
    "    - 若该节点为叶节点，将额外记录该叶节点的数值；\n",
    "- 若存在未到达的子节点，则选取未到达的一个子节点（优先选取左子节点）进入；\n",
    "- 若无子节点或所有子节点均到达过，此时记录 `9`，并回到当前节点的父节点（若存在）。\n",
    "\n",
    "如果最终记下的数依序连接成一个整数 `num`，满足 $num \\mod~p=target$，则视为解开谜题。\n",
    "请问有多少种二叉树的组装方案，可以使得最终记录下的数字可以解开谜题\n",
    "\n",
    "**注意：**\n",
    "- 两棵结构不同的二叉树，作为不同的组装方案\n",
    "- 两棵结构相同的二叉树且存在某个相同位置处的宝石编号不同，也作为不同的组装方案\n",
    "- 可能存在数值相同的两颗宝石\n",
    "\n",
    "**示例 1：**\n",
    "> 输入：`gem = [2,3]`\n",
    "> `p = 100000007`\n",
    "> `target = 11391299`\n",
    ">\n",
    "> 输出：`1`\n",
    ">\n",
    "> 解释：\n",
    "> 包含 `2` 个叶节点的结构只有一种。\n",
    "> 假设 B、C 节点的值分别为 3、2，对应 target 为 11391299，如下图所示。\n",
    "> 11391299 % 100000007 = 11391299，满足条件;\n",
    "> 假设 B、C 节点的值分别为 2、3，对应 target 为 11291399;\n",
    "> 11291399 % 100000007 = 11291399，不满足条件；\n",
    "> 因此只存在 1 种方案，返回 1\n",
    "![万灵 (1).gif](https://pic.leetcode.cn/1682397079-evMssw-%E4%B8%87%E7%81%B5%20\\(1\\).gif){:height=300px}\n",
    "\n",
    "\n",
    "**示例 2：**\n",
    "> 输入：`gem = [3,21,3]`\n",
    "> `p = 7`\n",
    "> `target = 5`\n",
    ">\n",
    "> 输出：`4`\n",
    ">\n",
    "> 解释：\n",
    "包含 `3` 个叶节点树结构有两种，列举如下：\n",
    "满足条件的组合有四种情况：\n",
    "> 当结构为下图（1）时：叶子节点的值为 [3,3,21] 或 [3,3,21]，得到的整数为 `11139139912199`。\n",
    "> 当结构为下图（2）时：叶子节点的值为 [21,3,3] 或 [21,3,3]，得到的整数为 `11219113913999`。\n",
    "![image.png](https://pic.leetcode.cn/1682322894-vfqJIV-image.png){:width=500px}\n",
    "\n",
    "\n",
    "**提示：**\n",
    "- `1 <= gem.length <= 9`\n",
    "- `0 <= gem[i] <= 10^9`\n",
    "- `1 <= p <= 10^9`，保证 $p$ 为素数。\n",
    "- `0 <= target < p`\n",
    "- 存在 2 组 `gem.length == 9` 的用例"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Description: [cnHoX6](https://leetcode.cn/problems/cnHoX6/description/)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Solutions: [cnHoX6](https://leetcode.cn/problems/cnHoX6/solutions/)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "test_cases = ['[2,3]\\n100000007\\n11391299', '[3,21,3]\\n7\\n5']"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def treeOfInfiniteSouls(self, gem: List[int], p: int, target: int) -> int:\n",
    "        N = len(gem)\n",
    "\n",
    "        # 先计算每个长度的余数\n",
    "        l_m = [1]\n",
    "\n",
    "        # 最大12 * 9 + 2 * 8 = 124\n",
    "        for _ in range(130):\n",
    "            l_m.append(l_m[-1] * 10 % p)\n",
    "\n",
    "        # 找到a，使ma = 1 mod p\n",
    "        if p > 10 and N > 6:\n",
    "            l_p = []\n",
    "            n = p - 2\n",
    "\n",
    "            while n:\n",
    "                n, m = divmod(n, 2)\n",
    "                l_p.append(m)\n",
    "\n",
    "            pn = len(l_p)\n",
    "\n",
    "            l_ma = [1] * len(l_m)\n",
    "            l_mp = [1] * pn\n",
    "\n",
    "            for mi in range(1, len(l_m)):\n",
    "                m = l_m[mi]\n",
    "\n",
    "                # 由于m ** (p - 1) = 1 mod p\n",
    "                # 所以 a == m ** (p - 2) mod p\n",
    "                l_mp[0] = m\n",
    "\n",
    "                for i in range(1, pn):\n",
    "                    l_mp[i] = (l_mp[i - 1] ** 2) % p\n",
    "\n",
    "                a = 1\n",
    "                for b, m in zip(l_p, l_mp):\n",
    "                    if b:\n",
    "                        a = a * m % p\n",
    "                l_ma[mi] = a % p\n",
    "\n",
    "        # 0~n-1这n个数里面去0~n个\n",
    "        l_qu = defaultdict(lambda: defaultdict(list))\n",
    "        l_qu[0][0] = [[(), ()]]\n",
    "\n",
    "        for n in range(1, N + 1):\n",
    "            l_qu[n][0] = [[(), tuple(range(n))]]\n",
    "\n",
    "            for i in range(1, n):\n",
    "                # 不包含n - 1\n",
    "                for l, r in l_qu[n - 1][i]:\n",
    "                    l_qu[n][i].append([l, r + (n - 1,)])\n",
    "\n",
    "                # 包含n - 1\n",
    "                for l, r in l_qu[n - 1][i - 1]:\n",
    "                    l_qu[n][i].append([l + (n - 1,), r])\n",
    "\n",
    "            l_qu[n][n] = [[tuple(range(n)), ()]]\n",
    "\n",
    "        # 动态规划\n",
    "        # key: (gem1, ..., gemn)\n",
    "        # value: {}, key-余数，value-数量\n",
    "        dp = {}\n",
    "        dp_len = {}\n",
    "\n",
    "        # 计算1个\n",
    "        all_len = (N - 1) * 2\n",
    "        for i in range(N):\n",
    "            s = \"1\" + str(gem[i]) + \"9\"\n",
    "            dp_len[(i,)] = len(s)\n",
    "            all_len += len(s)\n",
    "            dp[(i,)] = {int(s) % p: 1}\n",
    "\n",
    "        if N == 1:\n",
    "            return dp[(0,)].get(target, 0)\n",
    "\n",
    "        # 计算后续\n",
    "        # 若p小于10，直接计算到第\n",
    "        # 若N <= 6，也直接计算到第\n",
    "        if p < 10 or N <= 6:\n",
    "            L = N + 1\n",
    "        else:\n",
    "            L = 5\n",
    "\n",
    "        for n in range(2, L):\n",
    "            for l, _ in l_qu[N][n]:\n",
    "                d = defaultdict(int)\n",
    "                dp[l] = d\n",
    "\n",
    "                # l再左右拆分，左右都不能为空\n",
    "                # 考虑对称性，使0在左边，到时左右和交换后右做都计算一遍\n",
    "                for i in range(n - 1):\n",
    "                    for ll, lr in l_qu[n - 1][i]:\n",
    "                        ll = (l[0],) + tuple(l[i + 1] for i in ll)\n",
    "                        lr = tuple(l[i + 1] for i in lr)\n",
    "                        ln = dp_len[ll]\n",
    "                        rn = dp_len[lr]\n",
    "\n",
    "                        for lm, lnum in dp[ll].items():\n",
    "                            for rm, rnum in dp[lr].items():\n",
    "                                # 前面加1，lm*10**(ln + 1), rm*10，后面加9\n",
    "                                m = (l_m[ln + rn + 1] + lm * l_m[rn + 1] + 10 * rm + 9) % p\n",
    "                                d[m] += lnum * rnum\n",
    "\n",
    "                                # 左右对称\n",
    "                                m = (l_m[ln + rn + 1] + rm * l_m[ln + 1] + 10 * lm + 9) % p\n",
    "                                d[m] += lnum * rnum\n",
    "\n",
    "                dp_len[l] = ln + rn + 2\n",
    "\n",
    "        # 若p较大，计算最后一个单独计算加快速度\n",
    "        all = tuple(range(N))\n",
    "\n",
    "        if p > 10 and N > 6:\n",
    "            # for l, _ in l_qu[N][L - 1]:\n",
    "            #     print(\"dp\", l, len(dp[l]))\n",
    "\n",
    "            # 长度为L及以上的还没有计算，先计算需要的值\n",
    "            dp_need = defaultdict(lambda: defaultdict(set))\n",
    "            dp_need[N][all] = {target}\n",
    "            dp_len[all] = all_len\n",
    "\n",
    "            for n in range(N, L, -1):\n",
    "                for l, s in dp_need[n].items():\n",
    "                    lnrn = dp_len[l]\n",
    "\n",
    "                    for i in range(n - 1, L - 1, -1):\n",
    "                        d = dp_need[i]\n",
    "\n",
    "                        for ll, lr in l_qu[n][i]:\n",
    "                            ll = tuple(l[i] for i in ll)\n",
    "                            lr = tuple(l[i] for i in lr)\n",
    "\n",
    "                            rn = dp_len[lr]\n",
    "                            ln = lnrn - rn - 2\n",
    "                            dp_len[ll] = ln\n",
    "\n",
    "                            for rm in dp[lr]:\n",
    "                                for tm in s:\n",
    "                                    # 前面加1，lm*10**(ln + 1), rm*10，后面加9\n",
    "                                    m = (tm - l_m[ln + rn + 1] - rm * 10 - 9) % p\n",
    "\n",
    "                                    # A = l_m[rn + 1]\n",
    "                                    # 求解 Ax mod p = m\n",
    "                                    lm = m * l_ma[rn + 1] % p\n",
    "                                    d[ll].add(lm)\n",
    "\n",
    "                                    # 左右交换\n",
    "                                    m = (tm - l_m[ln + rn + 1] - rm * l_m[ln + 1] - 9) % p\n",
    "                                    lm = m * l_ma[1] % p\n",
    "                                    d[ll].add(lm)\n",
    "\n",
    "            # for n, d in dp_need.items():\n",
    "            #     for l, s in d.items():\n",
    "            #         print(\"need\", n, l, len(s))\n",
    "\n",
    "            # 计算L后面的\n",
    "            for n in range(L, N + 1):\n",
    "                for l, _ in l_qu[N][n]:\n",
    "                    d = defaultdict(int)\n",
    "                    dp[l] = d\n",
    "                    s = dp_need[n][l]\n",
    "\n",
    "                    # l再左右拆分，左右都不能为空\n",
    "                    # 考虑对称性，使0在左边，到时左右和交换后右做都计算一遍\n",
    "                    for i in range(n - 1):\n",
    "                        for ll, lr in l_qu[n - 1][i]:\n",
    "                            ll = (l[0],) + tuple(l[i + 1] for i in ll)\n",
    "                            lr = tuple(l[i + 1] for i in lr)\n",
    "                            ln = dp_len[ll]\n",
    "                            rn = dp_len[lr]\n",
    "\n",
    "                            for lm, lnum in dp[ll].items():\n",
    "                                for rm, rnum in dp[lr].items():\n",
    "                                    # 前面加1，lm*10**(ln + 1), rm*10，后面加9\n",
    "                                    m = (l_m[ln + rn + 1] + lm * l_m[rn + 1] + 10 * rm + 9) % p\n",
    "                                    if m in s:\n",
    "                                        d[m] += lnum * rnum\n",
    "\n",
    "                                    # 左右对称\n",
    "                                    m = (l_m[ln + rn + 1] + rm * l_m[ln + 1] + 10 * lm + 9) % p\n",
    "                                    if m in s:\n",
    "                                        d[m] += lnum * rnum\n",
    "\n",
    "        # for n in range(4, N + 1):\n",
    "        #     for l, _ in l_qu[N][n]:\n",
    "        #         print(l, len(dp[l]))\n",
    "\n",
    "        return dp[all][target]\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def treeOfInfiniteSouls(self, gem: List[int], p: int, target: int) -> int:\n",
    "        n = len(gem)\n",
    "        \n",
    "        # 预处理，10^-L 次方的逆元\n",
    "        # 以及 10^L\n",
    "        ten = [1]\n",
    "        teninv = [1]\n",
    "        tmp = 10\n",
    "        for _ in range(120): \n",
    "            tmp %= p\n",
    "            ten.append(tmp)           \n",
    "            # python自带快速幂\n",
    "            teninv.append(pow(tmp, p-2, p))\n",
    "            tmp *= 10\n",
    "                        \n",
    "        # 预处理，每个单state对应的数字\n",
    "        # state 状态下，拼成的字符串长度\n",
    "        L = [0] * (1 << n)\n",
    "        # 初始化，单个数字\n",
    "        # 预处理\n",
    "        # 单独处理单个1的状态对应的数，即单节点\n",
    "        pos = {}\n",
    "        for i,num in enumerate(gem):\n",
    "            state = 1 << i\n",
    "            L[state] = len(str(num))\n",
    "            pos[state] = (ten[L[state]+1] + num*10+9) % p\n",
    "        \n",
    "        # 记录每个状态bit_count                                                                 \n",
    "        k = [0] * (1<<n)\n",
    "        for state in range(1,1 << n):\n",
    "            # 获取其中一个子集即可，也不用获取排列，因为就算顺序不同，总长度也不会变\n",
    "            l = (state - 1) & state\n",
    "            r = state ^ l\n",
    "            # 左右子树前拼一个1，后拼一个9，因此长度+2\n",
    "            L[state] = L[l] + L[r] + 2        \n",
    "            k[state] = state.bit_count()                               \n",
    "                           \n",
    "        # 预处理\n",
    "        # 每一个状态对应的子集以及其补集\n",
    "        subsets = [[] for _ in range(1 << n)]\n",
    "        for i in range(1 << n):\n",
    "            s = i\n",
    "            flag = True\n",
    "            while s:\n",
    "                # 跳过自身\n",
    "                if flag:\n",
    "                    flag = False\n",
    "                else:\n",
    "                    subsets[i].append((s,i^s))\n",
    "                s = (s - 1) & i\n",
    "        #for a,b in subsets[(1<<n)-1]:\n",
    "        #    print(bin(i),bin(a),bin(b))\n",
    "                  \n",
    "        cnt = [Counter() for _ in range(1<<n)]\n",
    "        # 当前state下mod p的 Counter        \n",
    "        def dp(state):\n",
    "            if cnt[state]:\n",
    "                return \n",
    "                \n",
    "            if state in pos:\n",
    "                cnt[state][pos[state]] += 1\n",
    "                return\n",
    "            \n",
    "            # 遍历所有子集与对应补集\n",
    "            for ls,rs in subsets[state]:\n",
    "                if not cnt[ls]:\n",
    "                    dp(ls) \n",
    "                \n",
    "                if not cnt[rs]:\n",
    "                    dp(rs) \n",
    "                                    \n",
    "                # 后缀pushup合并\n",
    "                for x,xt in cnt[ls].items():\n",
    "                    for y,yt in cnt[rs].items():\n",
    "                        num = (ten[L[ls]+L[rs]+1] + x*ten[L[rs]+1]+y*10+9) % p\n",
    "                        # 乘法原理\n",
    "                        cnt[state][num] += xt * yt            \n",
    "        # p 为 5 和 2 比较特殊，因为10任意次方及其逆元都为0，10^(lx+ly+1) + x*10^(ly+1) + y*10 + 9，针对该公式，前面mod p都没掉，只剩下9，因此，若9%p==target，x,y取任何值都可以。可以用组合数学来解决，也可以直接走dp解决\n",
    "        if p == 5 or p == 2:                                      \n",
    "            dp((1<<n)-1) \n",
    "            return cnt[(1<<n)-1][target]\n",
    "                                                                                                           \n",
    "        # 10^(lx+ly+1) + x*10^(ly+1) + y*10 + 9                                                     \n",
    "        # 状态压缩state子集，满足tgt的方案数   \n",
    "        # 折半搜索\n",
    "        @cache       \n",
    "        def dfs(state,t):            \n",
    "            # 只有一个节点，那必定是子节点\n",
    "            if state in pos:               \n",
    "                return int(pos[state]==t)             \n",
    "            \n",
    "            if k[state] <= 6:\n",
    "                dp(state)\n",
    "                return cnt[state][t]\n",
    "                                                                                       \n",
    "            res = 0\n",
    "            # 遍历子集\n",
    "            for ls,rs in subsets[state]: \n",
    "                # 子集数目更少的直接查找\n",
    "                # 存y找x\n",
    "                if k[ls] >= k[rs]:\n",
    "                    dp(rs)\n",
    "                    # x = (t - y*10 - 9  - 10^(lx+ly+1))*10^-(ly+1) (mod m)                \n",
    "                    for y,ty in cnt[rs].items():                        \n",
    "                        nt = ((t-(y*10+9+ten[L[ls]+L[rs]+1])%p+p)*teninv[L[rs]+1]) % p                        \n",
    "                        res += dfs(ls,nt)*ty\n",
    "                # 存x找y                \n",
    "                else:\n",
    "                    dp(ls)\n",
    "                    # y =(t - x*10^(ly+1) - 9 - 10^(lx+ly+1))*10^-1 (mod m)\n",
    "                    for x,tx in cnt[ls].items():\n",
    "                        nt =((t-(x*ten[L[rs]+1]+9+ten[L[ls]+L[rs]+1])%p+p)*teninv[1])%p\n",
    "                        res += dfs(rs,nt)*tx                                                                           \n",
    "            return res\n",
    "                \n",
    "        res = dfs((1<<n)-1,target)        \n",
    "        dfs.cache_clear()\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def treeOfInfiniteSouls(self, gem: List[int], p: int, target: int) -> int:\n",
    "        n = len(gem)\n",
    "        \n",
    "        # 预处理，10^-L 次方的逆元\n",
    "        # 以及 10^L\n",
    "        ten = [1]\n",
    "        teninv = [1]\n",
    "        tmp = 10\n",
    "        for _ in range(120): \n",
    "            tmp %= p\n",
    "            ten.append(tmp)           \n",
    "            # python自带快速幂\n",
    "            teninv.append(pow(tmp, p-2, p))\n",
    "            tmp *= 10\n",
    "                        \n",
    "        # 预处理，每个单state对应的数字\n",
    "        # state 状态下，拼成的字符串长度\n",
    "        L = [0] * (1 << n)\n",
    "        # 初始化，单个数字\n",
    "        # 预处理\n",
    "        # 单独处理单个1的状态对应的数，即单节点\n",
    "        pos = {}\n",
    "        for i,num in enumerate(gem):\n",
    "            state = 1 << i\n",
    "            L[state] = len(str(num))\n",
    "            pos[state] = (ten[L[state]+1] + num*10+9) % p\n",
    "        \n",
    "        # 记录每个状态bit_count                                                                 \n",
    "        k = [0] * (1<<n)\n",
    "        for state in range(1,1 << n):\n",
    "            # 获取其中一个子集即可，也不用获取排列，因为就算顺序不同，总长度也不会变\n",
    "            l = (state - 1) & state\n",
    "            r = state ^ l\n",
    "            # 左右子树前拼一个1，后拼一个9，因此长度+2\n",
    "            L[state] = L[l] + L[r] + 2        \n",
    "            k[state] = state.bit_count()                               \n",
    "                           \n",
    "        # 预处理\n",
    "        # 每一个状态对应的子集以及其补集\n",
    "        subsets = [[] for _ in range(1 << n)]\n",
    "        for i in range(1 << n):\n",
    "            s = i\n",
    "            flag = True\n",
    "            while s:\n",
    "                # 跳过自身\n",
    "                if flag:\n",
    "                    flag = False\n",
    "                else:\n",
    "                    subsets[i].append((s,i^s))\n",
    "                s = (s - 1) & i\n",
    "        #for a,b in subsets[(1<<n)-1]:\n",
    "        #    print(bin(i),bin(a),bin(b))\n",
    "                  \n",
    "        cnt = [Counter() for _ in range(1<<n)]\n",
    "        # 当前state下mod p的 Counter        \n",
    "        def dp(state):\n",
    "            if cnt[state]:\n",
    "                return \n",
    "                \n",
    "            if state in pos:\n",
    "                cnt[state][pos[state]] += 1\n",
    "                return\n",
    "            \n",
    "            # 遍历所有子集与对应补集\n",
    "            for ls,rs in subsets[state]:\n",
    "                if not cnt[ls]:\n",
    "                    dp(ls) \n",
    "                \n",
    "                if not cnt[rs]:\n",
    "                    dp(rs) \n",
    "                                    \n",
    "                # 后缀pushup合并\n",
    "                for x,xt in cnt[ls].items():\n",
    "                    for y,yt in cnt[rs].items():\n",
    "                        num = (ten[L[ls]+L[rs]+1] + x*ten[L[rs]+1]+y*10+9) % p\n",
    "                        # 乘法原理\n",
    "                        cnt[state][num] += xt * yt            \n",
    "        # p 为 5 和 2 比较特殊，因为10任意次方及其逆元都为0，10^(lx+ly+1) + x*10^(ly+1) + y*10 + 9，针对该公式，前面mod p都没掉，只剩下9，因此，若9%p==target，x,y取任何值都可以。可以用组合数学来解决，也可以直接走dp解决\n",
    "        if p == 5 or p == 2:                                      \n",
    "            dp((1<<n)-1) \n",
    "            return cnt[(1<<n)-1][target]\n",
    "                                                                                                           \n",
    "        # 10^(lx+ly+1) + x*10^(ly+1) + y*10 + 9                                                     \n",
    "        # 状态压缩state子集，满足tgt的方案数   \n",
    "        # 折半搜索\n",
    "        @cache       \n",
    "        def dfs(state,t):            \n",
    "            # 只有一个节点，那必定是子节点\n",
    "            if state in pos:               \n",
    "                return int(pos[state]==t)             \n",
    "            \n",
    "            if k[state] <= 6:\n",
    "                dp(state)\n",
    "                return cnt[state][t]\n",
    "                                                                                       \n",
    "            res = 0\n",
    "            # 遍历子集\n",
    "            for ls,rs in subsets[state]: \n",
    "                # 子集数目更少的直接查找\n",
    "                # 存y找x\n",
    "                if k[ls] >= k[rs]:\n",
    "                    dp(rs)\n",
    "                    # x = (t - y*10 - 9  - 10^(lx+ly+1))*10^-(ly+1) (mod m)                \n",
    "                    for y,ty in cnt[rs].items():                        \n",
    "                        nt = ((t-(y*10+9+ten[L[ls]+L[rs]+1])%p+p)*teninv[L[rs]+1]) % p                        \n",
    "                        res += dfs(ls,nt)*ty\n",
    "                # 存x找y                \n",
    "                else:\n",
    "                    dp(ls)\n",
    "                    # y =(t - x*10^(ly+1) - 9 - 10^(lx+ly+1))*10^-1 (mod m)\n",
    "                    for x,tx in cnt[ls].items():\n",
    "                        nt =((t-(x*ten[L[rs]+1]+9+ten[L[ls]+L[rs]+1])%p+p)*teninv[1])%p\n",
    "                        res += dfs(rs,nt)*tx                                                                           \n",
    "            return res\n",
    "                \n",
    "        res = dfs((1<<n)-1,target)        \n",
    "        dfs.cache_clear()\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def treeOfInfiniteSouls(self, gem: List[int], p: int, target: int) -> int:\n",
    "        n = len(gem)\n",
    "        \n",
    "        # 预处理，10^-L 次方的逆元\n",
    "        # 以及 10^L\n",
    "        ten = [1]\n",
    "        teninv = [1]\n",
    "        tmp = 10\n",
    "        for _ in range(120): \n",
    "            tmp %= p\n",
    "            ten.append(tmp)           \n",
    "            # python自带快速幂\n",
    "            teninv.append(pow(tmp, p-2, p))\n",
    "            tmp *= 10\n",
    "                        \n",
    "        # 预处理，每个单state对应的数字\n",
    "        # state 状态下，拼成的字符串长度\n",
    "        L = [0] * (1 << n)\n",
    "        # 初始化，单个数字\n",
    "        # 预处理\n",
    "        # 单独处理单个1的状态对应的数，即单节点\n",
    "        pos = {}\n",
    "        for i,num in enumerate(gem):\n",
    "            state = 1 << i\n",
    "            L[state] = len(str(num))\n",
    "            pos[state] = (ten[L[state]+1] + num*10+9) % p\n",
    "        \n",
    "        # 记录每个状态bit_count                                                                 \n",
    "        k = [0] * (1<<n)\n",
    "        for state in range(1,1 << n):\n",
    "            # 获取其中一个子集即可，也不用获取排列，因为就算顺序不同，总长度也不会变\n",
    "            l = (state - 1) & state\n",
    "            r = state ^ l\n",
    "            # 左右子树前拼一个1，后拼一个9，因此长度+2\n",
    "            L[state] = L[l] + L[r] + 2        \n",
    "            k[state] = state.bit_count()                               \n",
    "                           \n",
    "        # 预处理\n",
    "        # 每一个状态对应的子集以及其补集\n",
    "        subsets = [[] for _ in range(1 << n)]\n",
    "        for i in range(1 << n):\n",
    "            s = i\n",
    "            flag = True\n",
    "            while s:\n",
    "                # 跳过自身\n",
    "                if flag:\n",
    "                    flag = False\n",
    "                else:\n",
    "                    subsets[i].append((s,i^s))\n",
    "                s = (s - 1) & i\n",
    "        #for a,b in subsets[(1<<n)-1]:\n",
    "        #    print(bin(i),bin(a),bin(b))\n",
    "                  \n",
    "        cnt = [Counter() for _ in range(1<<n)]\n",
    "        # 当前state下mod p的 Counter        \n",
    "        def dp(state):\n",
    "            if cnt[state]:\n",
    "                return \n",
    "                \n",
    "            if state in pos:\n",
    "                cnt[state][pos[state]] += 1\n",
    "                return\n",
    "            \n",
    "            # 遍历所有子集与对应补集\n",
    "            for ls,rs in subsets[state]:\n",
    "                if not cnt[ls]:\n",
    "                    dp(ls) \n",
    "                \n",
    "                if not cnt[rs]:\n",
    "                    dp(rs) \n",
    "                                    \n",
    "                # 后缀pushup合并\n",
    "                for x,xt in cnt[ls].items():\n",
    "                    for y,yt in cnt[rs].items():\n",
    "                        num = (ten[L[ls]+L[rs]+1] + x*ten[L[rs]+1]+y*10+9) % p\n",
    "                        # 乘法原理\n",
    "                        cnt[state][num] += xt * yt            \n",
    "        # p 为 5 和 2 比较特殊，因为10任意次方及其逆元都为0，10^(lx+ly+1) + x*10^(ly+1) + y*10 + 9，针对该公式，前面mod p都没掉，只剩下9，因此，若9%p==target，x,y取任何值都可以。可以用组合数学来解决，也可以直接走dp解决\n",
    "        if p == 5 or p == 2:                                      \n",
    "            dp((1<<n)-1) \n",
    "            return cnt[(1<<n)-1][target]\n",
    "                                                                                                           \n",
    "        # 10^(lx+ly+1) + x*10^(ly+1) + y*10 + 9                                                     \n",
    "        # 状态压缩state子集，满足tgt的方案数   \n",
    "        # 折半搜索\n",
    "        @cache       \n",
    "        def dfs(state,t):            \n",
    "            # 只有一个节点，那必定是子节点\n",
    "            if state in pos:               \n",
    "                return int(pos[state]==t)             \n",
    "            \n",
    "            if k[state] <= 6:\n",
    "                dp(state)\n",
    "                return cnt[state][t]\n",
    "                                                                                       \n",
    "            res = 0\n",
    "            # 遍历子集\n",
    "            for ls,rs in subsets[state]: \n",
    "                # 子集数目更少的直接查找\n",
    "                # 存y找x\n",
    "                if k[ls] >= k[rs]:\n",
    "                    dp(rs)\n",
    "                    # x = (t - y*10 - 9  - 10^(lx+ly+1))*10^-(ly+1) (mod m)                \n",
    "                    for y,ty in cnt[rs].items():                        \n",
    "                        nt = ((t-(y*10+9+ten[L[ls]+L[rs]+1])%p+p)*teninv[L[rs]+1]) % p                        \n",
    "                        res += dfs(ls,nt)*ty\n",
    "                # 存x找y                \n",
    "                else:\n",
    "                    dp(ls)\n",
    "                    # y =(t - x*10^(ly+1) - 9 - 10^(lx+ly+1))*10^-1 (mod m)\n",
    "                    for x,tx in cnt[ls].items():\n",
    "                        nt =((t-(x*ten[L[rs]+1]+9+ten[L[ls]+L[rs]+1])%p+p)*teninv[1])%p\n",
    "                        res += dfs(rs,nt)*tx                                                                           \n",
    "            return res\n",
    "                \n",
    "        res = dfs((1<<n)-1,target)        \n",
    "        dfs.cache_clear()\n",
    "        return res\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "class Solution:\n",
    "    def treeOfInfiniteSouls(self, gem: List[int], p: int, target: int) -> int:\n",
    "        n = len(gem)\n",
    "        \n",
    "        # 预处理，10^-L 次方的逆元\n",
    "        # 以及 10^L\n",
    "        ten = [1]\n",
    "        teninv = [1]\n",
    "        tmp = 10\n",
    "        for _ in range(120): \n",
    "            tmp %= p\n",
    "            ten.append(tmp)           \n",
    "            # python自带快速幂\n",
    "            teninv.append(pow(tmp, p-2, p))\n",
    "            tmp *= 10\n",
    "                        \n",
    "        # 预处理，每个单state对应的数字\n",
    "        # state 状态下，拼成的字符串长度\n",
    "        L = [0] * (1 << n)\n",
    "        # 初始化，单个数字\n",
    "        # 预处理\n",
    "        # 单独处理单个1的状态对应的数，即单节点\n",
    "        pos = {}\n",
    "        for i,num in enumerate(gem):\n",
    "            state = 1 << i\n",
    "            L[state] = len(str(num))\n",
    "            pos[state] = (ten[L[state]+1] + num*10+9) % p\n",
    "        \n",
    "        # 记录每个状态bit_count                                                                 \n",
    "        k = [0] * (1<<n)\n",
    "        for state in range(1,1 << n):\n",
    "            # 获取其中一个子集即可，也不用获取排列，因为就算顺序不同，总长度也不会变\n",
    "            l = (state - 1) & state\n",
    "            r = state ^ l\n",
    "            # 左右子树前拼一个1，后拼一个9，因此长度+2\n",
    "            L[state] = L[l] + L[r] + 2        \n",
    "            k[state] = state.bit_count()                               \n",
    "                           \n",
    "        # 预处理\n",
    "        # 每一个状态对应的子集以及其补集\n",
    "        subsets = [[] for _ in range(1 << n)]\n",
    "        for i in range(1 << n):\n",
    "            s = i\n",
    "            flag = True\n",
    "            while s:\n",
    "                # 跳过自身\n",
    "                if flag:\n",
    "                    flag = False\n",
    "                else:\n",
    "                    subsets[i].append((s,i^s))\n",
    "                s = (s - 1) & i\n",
    "        #for a,b in subsets[(1<<n)-1]:\n",
    "        #    print(bin(i),bin(a),bin(b))\n",
    "                  \n",
    "        cnt = [Counter() for _ in range(1<<n)]\n",
    "        # 当前state下mod p的 Counter        \n",
    "        def dp(state):\n",
    "            if cnt[state]:\n",
    "                return \n",
    "                \n",
    "            if state in pos:\n",
    "                cnt[state][pos[state]] += 1\n",
    "                return\n",
    "            \n",
    "            # 遍历所有子集与对应补集\n",
    "            for ls,rs in subsets[state]:\n",
    "                if not cnt[ls]:\n",
    "                    dp(ls) \n",
    "                \n",
    "                if not cnt[rs]:\n",
    "                    dp(rs) \n",
    "                                    \n",
    "                # 后缀pushup合并\n",
    "                for x,xt in cnt[ls].items():\n",
    "                    for y,yt in cnt[rs].items():\n",
    "                        num = (ten[L[ls]+L[rs]+1] + x*ten[L[rs]+1]+y*10+9) % p\n",
    "                        # 乘法原理\n",
    "                        cnt[state][num] += xt * yt            \n",
    "        # p 为 5 和 2 比较特殊，因为10任意次方及其逆元都为0，10^(lx+ly+1) + x*10^(ly+1) + y*10 + 9，针对该公式，前面mod p都没掉，只剩下9，因此，若9%p==target，x,y取任何值都可以。可以用组合数学来解决，也可以直接走dp解决\n",
    "        if p == 5 or p == 2:                                      \n",
    "            dp((1<<n)-1) \n",
    "            return cnt[(1<<n)-1][target]\n",
    "                                                                                                           \n",
    "        # 10^(lx+ly+1) + x*10^(ly+1) + y*10 + 9                                                     \n",
    "        # 状态压缩state子集，满足tgt的方案数   \n",
    "        # 折半搜索\n",
    "        @cache       \n",
    "        def dfs(state,t):            \n",
    "            # 只有一个节点，那必定是子节点\n",
    "            if state in pos:               \n",
    "                return int(pos[state]==t)             \n",
    "            \n",
    "            if k[state] <= 6:\n",
    "                dp(state)\n",
    "                return cnt[state][t]\n",
    "                                                                                       \n",
    "            res = 0\n",
    "            # 遍历子集\n",
    "            for ls,rs in subsets[state]: \n",
    "                # 子集数目更少的直接查找\n",
    "                # 存y找x\n",
    "                if k[ls] >= k[rs]:\n",
    "                    dp(rs)\n",
    "                    # x = (t - y*10 - 9  - 10^(lx+ly+1))*10^-(ly+1) (mod m)                \n",
    "                    for y,ty in cnt[rs].items():                        \n",
    "                        nt = ((t-(y*10+9+ten[L[ls]+L[rs]+1])%p+p)*teninv[L[rs]+1]) % p                        \n",
    "                        res += dfs(ls,nt)*ty\n",
    "                # 存x找y                \n",
    "                else:\n",
    "                    dp(ls)\n",
    "                    # y =(t - x*10^(ly+1) - 9 - 10^(lx+ly+1))*10^-1 (mod m)\n",
    "                    for x,tx in cnt[ls].items():\n",
    "                        nt =((t-(x*ten[L[rs]+1]+9+ten[L[ls]+L[rs]+1])%p+p)*teninv[1])%p\n",
    "                        res += dfs(rs,nt)*tx                                                                           \n",
    "            return res\n",
    "                \n",
    "        res = dfs((1<<n)-1,target)        \n",
    "        dfs.cache_clear()\n",
    "        return res\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from functools import cache\n",
    "from collections import Counter\n",
    "from math import perm\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def treeOfInfiniteSouls(self, gem: List[int], p: int, target: int) -> int:\n",
    "        if p in (2, 5):\n",
    "            if target == 9 % p:\n",
    "                return perm(len(gem) * 2 - 2, len(gem) - 1)\n",
    "            else:\n",
    "                return 0\n",
    "        LIMIT = len(gem) - len(gem) // 3\n",
    "        r10 = pow(10, p - 2, p)\n",
    "        def iter_mask(mask):\n",
    "            current = 0\n",
    "            while True:\n",
    "                current = (((current | ~mask) + 1) & mask)\n",
    "                if current == mask:\n",
    "                    break\n",
    "                else:\n",
    "                    yield current, (~current & mask)\n",
    "\n",
    "        @cache\n",
    "        def get_length(mask):\n",
    "            size = -2\n",
    "            for i, g in enumerate(gem):\n",
    "                if (1 << i) & mask:\n",
    "                    size += 4 + len(str(g))\n",
    "            return size\n",
    "\n",
    "        @cache\n",
    "        def find_all(mask):\n",
    "            if mask.bit_count() == 1:\n",
    "                return [int('1' + str(gem[mask.bit_length() - 1]) + '9') % p]\n",
    "            result = []\n",
    "            for left, right in iter_mask(mask):\n",
    "                left_length = get_length(left)\n",
    "                right_length = get_length(right)\n",
    "                base = pow(10, left_length + right_length + 1, p)\n",
    "                mul = pow(10, right_length + 1, p)\n",
    "                for lr in find_all(left):\n",
    "                    lm = lr * mul + base + 9\n",
    "                    for rr in find_all(right):\n",
    "                        result.append((lm + rr * 10) % p)\n",
    "            return result\n",
    "\n",
    "        @cache\n",
    "        def find_all_counter(mask):\n",
    "            return Counter(find_all(mask))\n",
    "\n",
    "        def find_target(mask, target_list):\n",
    "            if mask.bit_count() <= LIMIT:\n",
    "                return sum(find_all_counter(mask).get(t, 0) for t in target_list)\n",
    "            count = 0\n",
    "            for left, right in iter_mask(mask):\n",
    "                left_length = get_length(left)\n",
    "                right_length = get_length(right)\n",
    "                base = pow(10, left_length + right_length + 1, p)\n",
    "                if left.bit_count() <= right.bit_count():\n",
    "                    mul = pow(10, right_length + 1, p)\n",
    "                    rr_list = []\n",
    "                    for lr in find_all(left):\n",
    "                        lm = lr * mul + base + 9\n",
    "                        rr_list.extend(((t - lm) * r10) % p\n",
    "                                       for t in target_list)\n",
    "                    count += find_target(right, rr_list)\n",
    "                else:\n",
    "                    mulr = pow(r10, right_length + 1, p)\n",
    "                    lr_list = []\n",
    "                    for rr in find_all(right):\n",
    "                        rm = base + 9 + rr * 10\n",
    "                        lr_list.extend(((t - rm) * mulr) % p\n",
    "                                        for t in target_list)\n",
    "                    count += find_target(left, lr_list)\n",
    "            return count\n",
    "\n",
    "        return find_target((1 << len(gem)) - 1, [target])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from functools import cache\n",
    "from collections import Counter\n",
    "from math import perm\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def treeOfInfiniteSouls(self, gem: List[int], p: int, target: int) -> int:\n",
    "        if p in (2, 5):\n",
    "            if target == 9 % p:\n",
    "                return perm(len(gem) * 2 - 2, len(gem) - 1)\n",
    "            else:\n",
    "                return 0\n",
    "        LIMIT = len(gem) - len(gem) // 3\n",
    "        r10 = pow(10, p - 2, p)\n",
    "        def iter_mask(mask):\n",
    "            current = 0\n",
    "            while True:\n",
    "                current = (((current | ~mask) + 1) & mask)\n",
    "                if current == mask:\n",
    "                    break\n",
    "                else:\n",
    "                    yield current, (~current & mask)\n",
    "\n",
    "        @cache\n",
    "        def get_length(mask):\n",
    "            size = -2\n",
    "            for i, g in enumerate(gem):\n",
    "                if (1 << i) & mask:\n",
    "                    size += 4 + len(str(g))\n",
    "            return size\n",
    "\n",
    "        @cache\n",
    "        def find_all(mask):\n",
    "            if mask.bit_count() == 1:\n",
    "                return [int('1' + str(gem[mask.bit_length() - 1]) + '9') % p]\n",
    "            result = []\n",
    "            for left, right in iter_mask(mask):\n",
    "                left_length = get_length(left)\n",
    "                right_length = get_length(right)\n",
    "                base = pow(10, left_length + right_length + 1, p)\n",
    "                mul = pow(10, right_length + 1, p)\n",
    "                for lr in find_all(left):\n",
    "                    lm = lr * mul + base + 9\n",
    "                    for rr in find_all(right):\n",
    "                        result.append((lm + rr * 10) % p)\n",
    "            return result\n",
    "\n",
    "        @cache\n",
    "        def find_all_counter(mask):\n",
    "            return Counter(find_all(mask))\n",
    "\n",
    "        def find_target(mask, target_list):\n",
    "            if mask.bit_count() <= LIMIT:\n",
    "                return sum(find_all_counter(mask).get(t, 0) for t in target_list)\n",
    "            count = 0\n",
    "            for left, right in iter_mask(mask):\n",
    "                left_length = get_length(left)\n",
    "                right_length = get_length(right)\n",
    "                base = pow(10, left_length + right_length + 1, p)\n",
    "                if left.bit_count() <= right.bit_count():\n",
    "                    mul = pow(10, right_length + 1, p)\n",
    "                    rr_list = []\n",
    "                    for lr in find_all(left):\n",
    "                        lm = lr * mul + base + 9\n",
    "                        rr_list.extend(((t - lm) * r10) % p\n",
    "                                       for t in target_list)\n",
    "                    count += find_target(right, rr_list)\n",
    "                else:\n",
    "                    mulr = pow(r10, right_length + 1, p)\n",
    "                    lr_list = []\n",
    "                    for rr in find_all(right):\n",
    "                        rm = base + 9 + rr * 10\n",
    "                        lr_list.extend(((t - rm) * mulr) % p\n",
    "                                        for t in target_list)\n",
    "                    count += find_target(left, lr_list)\n",
    "            return count\n",
    "\n",
    "        return find_target((1 << len(gem)) - 1, [target])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from functools import cache\n",
    "from collections import Counter\n",
    "from math import perm\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def treeOfInfiniteSouls(self, gem: List[int], p: int, target: int) -> int:\n",
    "        if p in (2, 5):\n",
    "            if target == 9 % p:\n",
    "                return perm(len(gem) * 2 - 2, len(gem) - 1)\n",
    "            else:\n",
    "                return 0\n",
    "        LIMIT = len(gem) - len(gem) // 3\n",
    "        r10 = pow(10, p - 2, p)\n",
    "        def iter_mask(mask):\n",
    "            current = 0\n",
    "            while True:\n",
    "                current = (((current | ~mask) + 1) & mask)\n",
    "                if current == mask:\n",
    "                    break\n",
    "                else:\n",
    "                    yield current, (~current & mask)\n",
    "\n",
    "        @cache\n",
    "        def get_length(mask):\n",
    "            size = -2\n",
    "            for i, g in enumerate(gem):\n",
    "                if (1 << i) & mask:\n",
    "                    size += 4 + len(str(g))\n",
    "            return size\n",
    "\n",
    "        @cache\n",
    "        def find_all(mask):\n",
    "            if mask.bit_count() == 1:\n",
    "                return [int('1' + str(gem[mask.bit_length() - 1]) + '9') % p]\n",
    "            result = []\n",
    "            for left, right in iter_mask(mask):\n",
    "                left_length = get_length(left)\n",
    "                right_length = get_length(right)\n",
    "                base = pow(10, left_length + right_length + 1, p)\n",
    "                mul = pow(10, right_length + 1, p)\n",
    "                for lr in find_all(left):\n",
    "                    lm = lr * mul + base + 9\n",
    "                    for rr in find_all(right):\n",
    "                        result.append((lm + rr * 10) % p)\n",
    "            return result\n",
    "\n",
    "        @cache\n",
    "        def find_all_counter(mask):\n",
    "            return Counter(find_all(mask))\n",
    "\n",
    "        def find_target(mask, target_list):\n",
    "            if mask.bit_count() <= LIMIT:\n",
    "                return sum(find_all_counter(mask).get(t, 0) for t in target_list)\n",
    "            count = 0\n",
    "            for left, right in iter_mask(mask):\n",
    "                left_length = get_length(left)\n",
    "                right_length = get_length(right)\n",
    "                base = pow(10, left_length + right_length + 1, p)\n",
    "                if left.bit_count() <= right.bit_count():\n",
    "                    mul = pow(10, right_length + 1, p)\n",
    "                    rr_list = []\n",
    "                    for lr in find_all(left):\n",
    "                        lm = lr * mul + base + 9\n",
    "                        rr_list.extend(((t - lm) * r10) % p\n",
    "                                       for t in target_list)\n",
    "                    count += find_target(right, rr_list)\n",
    "                else:\n",
    "                    mulr = pow(r10, right_length + 1, p)\n",
    "                    lr_list = []\n",
    "                    for rr in find_all(right):\n",
    "                        rm = base + 9 + rr * 10\n",
    "                        lr_list.extend(((t - rm) * mulr) % p\n",
    "                                        for t in target_list)\n",
    "                    count += find_target(left, lr_list)\n",
    "            return count\n",
    "\n",
    "        return find_target((1 << len(gem)) - 1, [target])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from functools import cache\n",
    "from collections import Counter\n",
    "from math import perm\n",
    "\n",
    "\n",
    "class Solution:\n",
    "    def treeOfInfiniteSouls(self, gem: List[int], p: int, target: int) -> int:\n",
    "        if p in (2, 5):\n",
    "            if target == 9 % p:\n",
    "                return perm(len(gem) * 2 - 2, len(gem) - 1)\n",
    "            else:\n",
    "                return 0\n",
    "        LIMIT = len(gem) - len(gem) // 3\n",
    "        r10 = pow(10, p - 2, p)\n",
    "        def iter_mask(mask):\n",
    "            current = 0\n",
    "            while True:\n",
    "                current = (((current | ~mask) + 1) & mask)\n",
    "                if current == mask:\n",
    "                    break\n",
    "                else:\n",
    "                    yield current, (~current & mask)\n",
    "\n",
    "        @cache\n",
    "        def get_length(mask):\n",
    "            size = -2\n",
    "            for i, g in enumerate(gem):\n",
    "                if (1 << i) & mask:\n",
    "                    size += 4 + len(str(g))\n",
    "            return size\n",
    "\n",
    "        @cache\n",
    "        def find_all(mask):\n",
    "            if mask.bit_count() == 1:\n",
    "                return [int('1' + str(gem[mask.bit_length() - 1]) + '9') % p]\n",
    "            result = []\n",
    "            for left, right in iter_mask(mask):\n",
    "                left_length = get_length(left)\n",
    "                right_length = get_length(right)\n",
    "                base = pow(10, left_length + right_length + 1, p)\n",
    "                mul = pow(10, right_length + 1, p)\n",
    "                for lr in find_all(left):\n",
    "                    lm = lr * mul + base + 9\n",
    "                    for rr in find_all(right):\n",
    "                        result.append((lm + rr * 10) % p)\n",
    "            return result\n",
    "\n",
    "        @cache\n",
    "        def find_all_counter(mask):\n",
    "            return Counter(find_all(mask))\n",
    "\n",
    "        def find_target(mask, target_list):\n",
    "            if mask.bit_count() <= LIMIT:\n",
    "                return sum(find_all_counter(mask).get(t, 0) for t in target_list)\n",
    "            count = 0\n",
    "            for left, right in iter_mask(mask):\n",
    "                left_length = get_length(left)\n",
    "                right_length = get_length(right)\n",
    "                base = pow(10, left_length + right_length + 1, p)\n",
    "                if left.bit_count() <= right.bit_count():\n",
    "                    mul = pow(10, right_length + 1, p)\n",
    "                    rr_list = []\n",
    "                    for lr in find_all(left):\n",
    "                        lm = lr * mul + base + 9\n",
    "                        rr_list.extend(((t - lm) * r10) % p\n",
    "                                       for t in target_list)\n",
    "                    count += find_target(right, rr_list)\n",
    "                else:\n",
    "                    mulr = pow(r10, right_length + 1, p)\n",
    "                    lr_list = []\n",
    "                    for rr in find_all(right):\n",
    "                        rm = base + 9 + rr * 10\n",
    "                        lr_list.extend(((t - rm) * mulr) % p\n",
    "                                        for t in target_list)\n",
    "                    count += find_target(left, lr_list)\n",
    "            return count\n",
    "\n",
    "        return find_target((1 << len(gem)) - 1, [target])\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from typing import List\n",
    "import collections\n",
    "\n",
    "from functools import cache\n",
    "from collections import Counter\n",
    "from math import perm\n",
    "from typing import List\n",
    "\n",
    "class Solution:\n",
    "    def treeOfInfiniteSouls(self, gem: List[int], p: int, target: int) -> int:\n",
    "        if p in (2, 5):\n",
    "            if target == 9 % p:\n",
    "                return perm(len(gem) * 2 - 2, len(gem) - 1)\n",
    "            else:\n",
    "                return 0\n",
    "        LIMIT = len(gem) - len(gem) // 3\n",
    "        r10 = pow(10, p - 2, p)\n",
    "        \n",
    "        def iter_mask(mask):\n",
    "            current = 0\n",
    "            while True:\n",
    "                current = (((current | ~mask) + 1) & mask)\n",
    "                if current == mask:\n",
    "                    break\n",
    "                else:\n",
    "                    yield current, (~current & mask)\n",
    "\n",
    "        @cache\n",
    "        def get_length(mask):\n",
    "            size = -2\n",
    "            for i, g in enumerate(gem):\n",
    "                if (1 << i) & mask:\n",
    "                    size += 4 + len(str(g))\n",
    "            return size\n",
    "\n",
    "        @cache\n",
    "        def find_all(mask):\n",
    "            if mask.bit_count() == 1:\n",
    "                return [int('1' + str(gem[mask.bit_length() - 1]) + '9') % p]\n",
    "            result = []\n",
    "            for left, right in iter_mask(mask):\n",
    "                left_length = get_length(left)\n",
    "                right_length = get_length(right)\n",
    "                base = pow(10, left_length + right_length + 1, p)\n",
    "                mul = pow(10, right_length + 1, p)\n",
    "                for lr in find_all(left):\n",
    "                    lm = lr * mul + base + 9\n",
    "                    for rr in find_all(right):\n",
    "                        result.append((lm + rr * 10) % p)\n",
    "            return result\n",
    "\n",
    "        @cache\n",
    "        def find_all_counter(mask):\n",
    "            return Counter(find_all(mask))\n",
    "\n",
    "        def find_target(mask, target_list):\n",
    "            if mask.bit_count() <= LIMIT:\n",
    "                return sum(find_all_counter(mask).get(t, 0) for t in target_list)\n",
    "            count = 0\n",
    "            for left, right in iter_mask(mask):\n",
    "                left_length = get_length(left)\n",
    "                right_length = get_length(right)\n",
    "                base = pow(10, left_length + right_length + 1, p)\n",
    "                if left.bit_count() <= right.bit_count():\n",
    "                    mul = pow(10, right_length + 1, p)\n",
    "                    rr_list = []\n",
    "                    for lr in find_all(left):\n",
    "                        lm = lr * mul + base + 9\n",
    "                        rr_list.extend(((t - lm) * r10) % p\n",
    "                                       for t in target_list)\n",
    "                    count += find_target(right, rr_list)\n",
    "                else:\n",
    "                    mulr = pow(r10, right_length + 1, p)\n",
    "                    lr_list = []\n",
    "                    for rr in find_all(right):\n",
    "                        rm = base + 9 + rr * 10\n",
    "                        lr_list.extend(((t - rm) * mulr) % p\n",
    "                                        for t in target_list)\n",
    "                    count += find_target(left, lr_list)\n",
    "            return count\n",
    "\n",
    "        return find_target((1 << len(gem)) - 1, [target])\n",
    "\n"
   ]
  }
 ],
 "metadata": {},
 "nbformat": 4,
 "nbformat_minor": 2
}
